#pragma once

#include <string>
#include <unordered_map>
#include <utility>

#include "utils/filesystem/filesystem.h"
#include "backend/kernel_interface/kernel_scheduler.h"

namespace txdnn {

class KernelCache {
public:
  using Key         = std::pair<fs::path, std::string>;
  using KernelMap   = std::unordered_map<Key, std::vector<backend::Kernel>, SimpleHash>;
  using ProgramMap  = std::unordered_map<Key, backend::Program, SimpleHash>;

  KernelCache() {}

  backend::Kernel addKernel(const Handle& h,
                            const std::string& algorithm,
                            const std::string& cfg,
                            const fs::path& program_path,
                            const std::string& kernel_name,
                            const std::vector<size_t>& vld, 
                            const std::vector<size_t>& vgd,
                            std::string params               = "",
                            std::size_t cache_idx            = 0,
                            const std::string& kernel_src    = "",
                            backend::Program* program_out    = nullptr) {

    const std::pair<std::string, std::string> key = std::make_pair(algorithm, cfg);
    if (!cfg.empty() || !algorithm.empty()) {
      // TXLOG_DEBUG();
    }

    const auto program = [&] {
      auto program_it = program_map_.find(std::make_pair(program_path, params));
      if (program_it != program_map_.end()) {
        auto& program = program_it->second;
        if (program_out != nullptr && !program.isCodeObjectInMemory() &&
            !program.isCodeObjectInFile()) {
          program = h.loadProgram(program_path, params, kernel_src, true);
        }
        return program;
      } else {
        auto program = h.loadProgram(program_path, params, kernel_src, program_out != nullptr);
        program_map_[std::make_pair(program_path, params)] = program;
        return program;
      }
    }();

    if (program_out != nullptr) {
      *program_out = program;
    }

    backend::Kernel kernel{};
    if (true) { // TODO check arch
      kernel = backend::Kernel{program, kernel_name};
    } else {
      kernel = backend::Kernel{program, kernel_name, vld, vgd};
    }

    if (!cfg.empty() && !algorithm.empty()) {
      addKernel(key, kernel, cache_idx);
    }
    return kernel;
  }

  void addKernel(Key key, backend::Kernel k, std::size_t cache_idx) {
    auto&& v = kernel_map_[key];
    if (cache_idx >= v.size()) {
      v.resize(cache_idx + 1);
    }
    v[cache_idx] = k;
  }

  void addProgram(backend::Program p, const fs::path& program_path, std::string params) {
    program_map_[std::make_pair(program_path, params)] = p;
  }

private:
  KernelMap kernel_map_;
  ProgramMap program_map_;
};


} // namespace txdnn